# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import tabulate

FLAGS = flags.FLAGS


def calc_edge(bid_x, ask_x, bid_y, ask_y):
  gap = bid_x - ask_y
  mid = (bid_x + ask_x + bid_y + ask_y) / 4
  edge_bp = gap / mid * 10000
  return edge_bp


def calc_arb_stats(df, exchange, product):
  rows = []
  headers = [
      'exchange_pair',
      'product',
      'edge_bp',
      'fluctuate_count',
      'Fcoin->another_count',
      'another->Fcoin_count'
  ]
  for edge_bp in [0, 2, 5, 10, 15, 20]:
    fluctuate_count = 0
    x_to_y_count = 0
    y_to_x_count = 0
    continuous_x_to_y_count = 0  # 10 times for a cycle
    continuous_y_to_x_count = 0  # 10 times for a cycle
    for _, row in df.iterrows():
      curr_edge_x_to_y = calc_edge(row['true_bid_x'],
                                   row['true_ask_x'],
                                   row['true_bid_y'],
                                   row['true_ask_y'])
      curr_edge_y_to_x = calc_edge(row['true_bid_y'],
                                   row['true_ask_y'],
                                   row['true_bid_x'],
                                   row['true_ask_x'])
      if curr_edge_x_to_y > edge_bp:
        x_to_y_count += 1
        continuous_x_to_y_count += 1
        continuous_y_to_x_count = 0
      elif curr_edge_y_to_x > edge_bp:
        y_to_x_count += 1
        continuous_y_to_x_count += 1
        continuous_x_to_y_count = 0

      if continuous_x_to_y_count == 10:
        continuous_x_to_y_count += 1
        fluctuate_count += 1
      elif continuous_y_to_x_count == 10:
        continuous_y_to_x_count += 1
        fluctuate_count += 1

    row = ['Fcoin/%s' % exchange, product, edge_bp, fluctuate_count, x_to_y_count, y_to_x_count]
    rows.append(row)
  print(tabulate.tabulate(rows, headers=headers))
  print()


def merge_dfs(df1, df2):
  return pandas.merge(df1, df2, how='inner', left_index=True, right_index=True)


def sample_df_per_second(df):
  return df.resample('S', closed='left', label='left').last()


def read_csv_into_df(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
  df = df.set_index('datetime')
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  assert csv_root, '--csv_root must be specified.'

  products = ['BTC-USDT', 'ETC-USDT', 'ETH-USDT']
  for product in products:
    fcoin_price_df = read_csv_into_df(csv_root, '%s.%s.csv' % (product, 'Fcoin'))
    for exchange in ['Binance', 'Bitfinex', 'Huobi', 'Okex']:
      other_price_df = read_csv_into_df(csv_root, '%s.%s.csv' % (product, exchange))
      sampled_df1 = sample_df_per_second(fcoin_price_df)
      sampled_df2 = sample_df_per_second(other_price_df)
      merged_df = merge_dfs(sampled_df1, sampled_df2)
      calc_arb_stats(merged_df, exchange, product)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Input csv files root directory.')

  app.run(main)
